#include <iostream>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/epoll.h>
#include <unistd.h>
#include <fcntl.h>
#include <cstring>
#include <cstdlib>
#include <cerrno>
#include <cassert>
#include <thread>

using namespace std;

const int MAX_EVENT_NUMBER = 1024;
const int BUFFER_SIZE = 1024;

int setNonBlocking(int fd) {
    int oldOption = fcntl(fd, F_GETFL);
    int newOption = oldOption | O_NONBLOCK;
    fcntl(fd, F_SETFL, newOption);
    return oldOption;
}

void addFd(int epollFd, int fd, bool oneShot = false) {
    epoll_event event;
    event.data.fd = fd;
    event.events = EPOLLIN | EPOLLET;
    if (oneShot) {
        event.events |= EPOLLONESHOT;   // 开启 EPOLLONESHOT 模式
    }
    epoll_ctl(epollFd, EPOLL_CTL_ADD, fd, &event);
    setNonBlocking(fd);
}

void resetOneShot(int epollFd, int fd) {
    epoll_event event;
    event.data.fd = fd;
    event.events = EPOLLIN | EPOLLET | EPOLLONESHOT; // 重置为 EPOLLONESHOT 模式
    epoll_ctl(epollFd, EPOLL_CTL_MOD, sockFd, &event);
}

void worker(int epollFd, int sockFd) {
    cout << "start new thread to receive data on fd: " << sockFd << endl;
    char buf[BUFFER_SIZE];
    memset(buf, 0, BUFFER_SIZE);
    while (true) {
        int ret = recv(sockFd, buf, BUFFER_SIZE - 1, 0);
        if (ret < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                cout << "read later" << endl;
                resetOneShot(epollFd, sockFd); // 重置为 EPOLLONESHOT 模式
                break;
            }
            close(sockFd);
            cout << "recv error: " << strerror(errno) << endl;
            break;
        } 
        else if (ret == 0) {
            close(sockFd);
            cout << "client closed connection" << endl;
            break;
        } 
        else {
            cout << "get " << ret << " bytes of content: " << buf << endl;
            this_thread::sleep_for(chrono::seconds(5)); // 模拟处理数据的时间
        }
    }
    cout << "end thread receiving data on fd: " << sockFd << endl;
}

int main(int argc, char* argv[])
{
    if(argv <= 2) {
        cerr << "Usage: " << argv[0] << " <IP> <Port>" << endl;
        return 1;
    }
    const char* ip = argv[1];
    int port = atoi(argv[2]);

    struct sockaddr_in address;
    address.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &address.sin_addr);
    address.sin_port = htons(port);

    int listenFd = socket(AF_INET, SOCK_STREAM, 0);
    assert(listenFd >= 0);

    int ret = bind(listenFd, (struct sockaddr*)&address, sizeof(address));
    assert(ret != -1);

    ret = listen(listenFd, 5);
    assert(ret != -1);

    int epollFd = epoll_create(5);
    assert(epollFd != -1);
    epoll_event events[MAX_EVENT_NUMBER];
    addFd(epollFd, listenFd, false);  // 监听套接字不能开启 EPOLLONESHOT 模式，后续的连接不会触发EPOLLIN事件

    while (true) {
        int ret = epoll_wait(epollFd, events, MAX_EVENT_NUMBER, -1);
        if (ret < 0 && errno != EINTR) {
            cerr << "epoll failure: " << strerror(errno) << endl;
            break;
        }
        
        for(int i = 0; i < ret; ++i) {
            int sockFd = events[i].data.fd;
            if (sockFd == listenFd) {
                struct sockaddr_in client;
                socklen_t clientLen = sizeof(client);
                int connFd = accept(listenFd, (struct sockaddr*)&client, &clientLen);
                addFd(epollFd, connFd, true); // 新连接开启 EPOLLONESHOT 模式
            } 
            else if (events[i].events & EPOLLIN) {
                thread(worker, epollFd, sockFd).detach(); // 分离线程，避免资源泄漏
            } 
            else {
                cout << "some other event occurred on fd: " << sockFd << endl;
            }
        }
        
    }

    close(listenFd);
    close(epollFd);
    return 0;
}